import torch
import torch.nn as nn

from models.custom import ResDDFN, BlockRes2ConvBR
from models.osmnet import L2NormDense, input_norm

class ImgPyramidModule(nn.Module):

    def __init__(self, out_ch=32):
        super().__init__()

        in_ch = 1


        self.Inc1 = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, 
                padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.LeakyReLU(inplace=True)
        )
        self.Inc2 = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, 
                stride=2, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.LeakyReLU(inplace=True), 
            nn.ConvTranspose2d(
                out_ch, out_ch, kernel_size=3, 
                stride=2, padding=1, bias=False, 
                output_padding=1)
        )

        self.transition = nn.Sequential(
            nn.Conv2d(out_ch*2, out_ch, kernel_size=3, 
                padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
        )

    def forward(self, x):
        ft1 = self.Inc1(x)
        ft2 = self.Inc2(x)
        ft = torch.cat([ft1, ft2], dim=1)

        out = self.transition(ft)

        return out


class ResDDFNIPM(ResDDFN):

    def __init__(self, out_ch=9):
        super().__init__(out_ch)

        self.conv1 = ImgPyramidModule(out_ch=32)


class ResBlock3(nn.Module):

    def __init__(self, out_ch=9):
        super().__init__()

        self.in_channels = 32

        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True)
        )

        self.conv2_x = BlockRes2ConvBR(32, 32)
        self.conv3_x = BlockRes2ConvBR(32, 64)
        self.conv4_x = BlockRes2ConvBR(64, 128)

        self.conv2 = nn.Sequential(
            nn.Dropout(0.1),
            nn.Conv2d(128, out_ch, kernel_size=1,
                bias=False),
            nn.BatchNorm2d(out_ch)
        )

    def forward(self, input):
        input = input_norm(input)

        input = self.conv1(input)
        input = self.conv2_x(input)
        input = self.conv3_x(input)
        sparse_ft = self.conv4_x(input)
        feature = self.conv2(sparse_ft)

        L2norm_feature = L2NormDense()(feature)
        return feature, L2norm_feature, sparse_ft


if __name__ == "__main__":

    # net = ImgPyramidModule().cuda()
    net = ResDDFNIPM().cuda()

    img = torch.zeros((2, 1, 128, 128)).float().cuda()

    _, out = net(img)
    print(out.shape)
